import os
import torch
from torchvision import datasets, transforms, models
from torch.utils.data import DataLoader

# -----------------------------
# 🟢 配置参数
# -----------------------------
MODEL_PATH = "cat_dog_model_mps.pth"  # 已训练好的模型文件
BATCH_SIZE = 32
VAL_DIR = "data/processed/val"        # 验证集路径

# -----------------------------
# 🟢 设备检测
# -----------------------------
device = torch.device("mps" if torch.backends.mps.is_available() else "cpu")
print(f"📌 当前设备: {device}")
print(f"🔍 PyTorch 版本: {torch.__version__}")
print(f"🔋 是否支持 MPS: {torch.backends.mps.is_available()}")
print(f"🔋 是否支持 CUDA: {torch.cuda.is_available()}")

# -----------------------------
# 🟢 数据预处理
# -----------------------------
transform = transforms.Compose([
    transforms.Resize((128, 128)),
    transforms.ToTensor(),
])

val_dataset = datasets.ImageFolder(VAL_DIR, transform=transform)
val_loader = DataLoader(val_dataset, batch_size=BATCH_SIZE, shuffle=False, num_workers=0)
print(f"📦 验证样本数量: {len(val_dataset)}")

# -----------------------------
# 🟢 模型加载
# -----------------------------
model = models.resnet18(weights=None)  # 不使用预训练权重
model.fc = torch.nn.Linear(model.fc.in_features, 2)  # 2 类（猫/狗）
model.to(device)

# 加载已有训练模型
if os.path.exists(MODEL_PATH):
    model.load_state_dict(torch.load(MODEL_PATH, map_location=device))
    print(f"✅ 模型已加载: {MODEL_PATH}")
else:
    raise FileNotFoundError(f"模型文件不存在: {MODEL_PATH}")

# -----------------------------
# 🟢 验证逻辑
# -----------------------------
model.eval()
correct = 0
total = 0

with torch.no_grad():
    for images, labels in val_loader:
        images, labels = images.to(device), labels.to(device)
        outputs = model(images)
        _, predicted = torch.max(outputs.data, 1)
        total += labels.size(0)
        correct += (predicted == labels).sum().item()

accuracy = 100 * correct / total
print(f"🎯 验证集准确率: {accuracy:.2f}%")
